import * as vscode from 'vscode';
import { ProxyService } from "../../../src/services/proxyService";
import { VsCodeWebviewProtocol } from "../../../src/webviewProtocol";
import { AnalyticsService } from '../../../src/services/analyticsService';

interface ToolCallChunk {
    index: number;
    id?: string;
    type?: 'function';
    function?: {
        name?: string;
        arguments?: string;
    };
}

export class AiChatWebviewRequestHandler {
    private proxyService: ProxyService;
    private currentAbortController: AbortController | null = null;

    constructor(private context: vscode.ExtensionContext, private webviewProtocol: VsCodeWebviewProtocol) {
        this.context = context;
        this.proxyService = ProxyService.getInstance(this.context);
        this.webviewProtocol = webviewProtocol;
    }

    /**
     * 处理聊天请求
     * @param message 
     */
    public async chatRequest(message: any){
        try {
            // 创建新的 AbortController
            this.currentAbortController = new AbortController();
            
            AnalyticsService.getInstance(this.context).trackEvent('aiChat');
            
            const requestData = JSON.parse(message.data);
            
            // 使用支持 AbortSignal 的 streamChat 方法
            const stream = this.proxyService.streamChat(requestData, this.currentAbortController.signal);
            
            const toolCalls: ToolCallChunk[] = [];
            let hasToolCall = false;
    
            for await (const delta of stream) {
                // 检查是否已被取消
                if (this.currentAbortController?.signal.aborted) {
                    break;
                }
                
                if (delta?.content) {
                    this.webviewProtocol.send('streamChunk', { content: delta.content });
                }
                
                if (delta?.tool_calls) {
                    hasToolCall = true;
                    delta.tool_calls.forEach((toolCallChunk: ToolCallChunk) => {
                        const existing = toolCalls.find(tc => tc.index === toolCallChunk.index);
                        if (existing) {
                            if (toolCallChunk.id) { existing.id = toolCallChunk.id; }
                            if (toolCallChunk.type) { existing.type = toolCallChunk.type; }
                            if (toolCallChunk.function) {
                              if (!existing.function) {
                                  existing.function = {name: "", arguments: ""};
                              }
                              if (toolCallChunk.function.name) {
                                  existing.function.name = (existing.function.name || '') + toolCallChunk.function.name;
                              }
                              if (toolCallChunk.function.arguments) {
                                  existing.function.arguments = (existing.function.arguments || '') + toolCallChunk.function.arguments;
                              }
                            }
                        } else {
                            toolCalls.push(JSON.parse(JSON.stringify(toolCallChunk)));
                        }
                    });
                }
            }
    
            if (hasToolCall) {
                const finalToolCalls = toolCalls
                  .filter(tc => tc.id && tc.function && tc.function.name && tc.function.arguments)
                  .map(tc => ({
                    id: tc.id!,
                    type: 'function',
                    function: {
                        name: tc.function!.name!,
                        arguments: tc.function!.arguments!,
                    },
                }));
                this.webviewProtocol.send('streamToolCalls', { toolCalls: finalToolCalls });
            }
            
            if (!this.currentAbortController?.signal.aborted) {
                this.webviewProtocol.send('streamEnd', {});
            }
    
          } catch (error: any) {
            console.error('Failed to start chat stream:', error);
            if (error.message === 'Request aborted' || this.currentAbortController?.signal.aborted) {
                this.webviewProtocol.send('streamError', { error: '聊天请求已被取消' });
            } else {
                this.webviewProtocol.send('streamError', { error: error.message });
            }
          } finally {
            // 清理 AbortController
            this.currentAbortController = null;
          }
    }

    /**
     * 处理停止聊天请求
     */
    public stopChatRequest(): void {
        if (this.currentAbortController) {
            this.currentAbortController.abort();
            this.currentAbortController = null;
            console.log('Chat request stopped by user');
        }
    }

}